/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import Unsafe from './Unsafe';
import HuffmanCompressor from './HuffmanCompressor';
import HuffmanCompressionTable from './HuffmanCompressionTable';
import Histogram from './Histogram';
import Huffman from './Huffman';
import SequenceEncoder from './SequenceEncoder';
import { CompressionParameters, Strategy } from './CompressionParameters';
import XxHash64 from './XxHash64';
import Long from '../util/long/index'
import Constants from './Constants'
import CompressionContext from './CompressionContext'
import HuffmanCompressionContext from './HuffmanCompressionContext'
import NumberTransform from './NumberTransform'
import Util from './Util'
import Exception from '../util/Exception'

export default class ZstdFrameCompressor {
    static MAX_FRAME_HEADER_SIZE: number = 14;
    private static CHECKSUM_FLAG: number = 0b100;
    private static SINGLE_SEGMENT_FLAG: number = 0b100000;
    private static MINIMUM_LITERALS_SIZE: number = 63;
    private static MAX_HUFFMAN_TABLE_LOG: number = 11;

    constructor() {
    }

    static writeMagic(outputBase: any, outputAddress: Long, outputLimit: Long): number
    {
        Util.checkArgument(
        outputLimit.sub(outputAddress)
            .greaterThanOrEqual(Constants.SIZE_OF_INT), "Output buffer too small");

        Unsafe.putInt(outputBase, outputAddress.toNumber(), Constants.MAGIC_NUMBER);
        return Constants.SIZE_OF_INT;
    }

    public static highestOneBit(i: number) {
        // HD, Figure 3-1
        i |= (i >> 1);
        i |= (i >> 2);
        i |= (i >> 4);
        i |= (i >> 8);
        i |= (i >> 16);
        return i - (i >>> 1);
    }

    public static numberOfLeadingZeros(i: number) {
        // HD, Figure 5-6
        if (i == 0)
        return 32;
        let n = 1;
        if (i >>> 16 == 0) {
            n += 16;
            i <<= 16;
        }
        if (i >>> 24 == 0) {
            n += 8;
            i <<= 8;
        }
        if (i >>> 28 == 0) {
            n += 4;
            i <<= 4;
        }
        if (i >>> 30 == 0) {
            n += 2;
            i <<= 2;
        }
        n -= i >>> 31;
        return n;
    }

    static writeFrameHeader(outputBase: any, outputAddress: Long, outputLimit: Long, inputSize: number, windowSize: number): number
    {
        Util.checkArgument(
        outputLimit.sub(outputAddress)
            .greaterThanOrEqual(ZstdFrameCompressor.MAX_FRAME_HEADER_SIZE), "Output buffer too small");

        let output: Long = outputAddress;

        let contentSizeDescriptor: number = (inputSize >= 256 ? 1 : 0) + (inputSize >= 65536 + 256 ? 1 : 0);
        let frameHeaderDescriptor: number = (contentSizeDescriptor << 6) | ZstdFrameCompressor.CHECKSUM_FLAG; // dictionary ID missing

        let singleSegment: boolean = windowSize >= inputSize;
        if (singleSegment) {
            frameHeaderDescriptor |= ZstdFrameCompressor.SINGLE_SEGMENT_FLAG;
        }

        Unsafe.putByte(outputBase, output.toNumber(), NumberTransform.toByte(frameHeaderDescriptor));
        output = output.add(1);

        if (!singleSegment) {
            let base: number = ZstdFrameCompressor.highestOneBit(windowSize);

            let exponent: number = 32 - ZstdFrameCompressor.numberOfLeadingZeros(base) - 1;
            if (exponent < Constants.MIN_WINDOW_LOG) {
                throw new Exception("Minimum window size is " + (1 << Constants.MIN_WINDOW_LOG));
            }

            let remainder: number = windowSize - base;
            if (remainder % (base / 8) != 0) {
                throw new Exception("Window size of magnitude 2^" + exponent + " must be multiple of " + (base / 8));
            }

            let mantissa: number = remainder / (base / 8);
            let encoded: number = ((exponent - Constants.MIN_WINDOW_LOG) << 3) | mantissa;

            Unsafe.putByte(outputBase, output.toNumber(), encoded);
            output = output.add(1);
        }

        switch (contentSizeDescriptor) {
            case 0:
                if (singleSegment) {
                    Unsafe.putByte(outputBase, output.toNumber(), inputSize);
                    output = output.add(1)
                }
                break;
            case 1:
                Unsafe.putShort(outputBase, output.toNumber(), inputSize - 256);
                output = output.add(Constants.SIZE_OF_SHORT);
                break;
            case 2:
                Unsafe.putInt(outputBase, output.toNumber(), inputSize);
                output = output.add(Constants.SIZE_OF_INT)
                break;
            default:
                throw new Error();
        }

        return output.sub(outputAddress).toInt();
    }

    static writeChecksum(outputBase: any, outputAddress: Long, outputLimit: Long, inputBase: Object, inputAddress: Long, inputLimit: Long): number
    {
        Util.checkArgument(
        outputLimit.sub(outputAddress)
            .greaterThanOrEqual(Constants.SIZE_OF_INT), "Output buffer too small");

        let inputSize: number = inputLimit.sub(inputAddress).toInt();

        let hash: Long = XxHash64.hash(Long.fromNumber(0), inputBase, inputAddress, inputSize);

        Unsafe.putInt(outputBase, outputAddress.toNumber(), hash.toInt());

        return Constants.SIZE_OF_INT;
    }

    public static compress(inputBase: Object, inputAddress: Long, inputLimit: Long, outputBase: Object, outputAddress: Long, outputLimit: Long, compressionLevel: number): number
    {
        let inputSize: number = inputLimit.sub(inputAddress).toInt();

        let parameters: CompressionParameters = CompressionParameters.compute(compressionLevel, inputSize);

        let output: Long = outputAddress;

        output = output.add(ZstdFrameCompressor.writeMagic(outputBase, output, outputLimit));
        output = output.add(ZstdFrameCompressor.writeFrameHeader(outputBase, output, outputLimit, inputSize, 1 << parameters.getWindowLog()));
        output = output.add(ZstdFrameCompressor.compressFrame(inputBase, inputAddress, inputLimit, outputBase, output, outputLimit, parameters));
        output = output.add(ZstdFrameCompressor.writeChecksum(outputBase, output, outputLimit, inputBase, inputAddress, inputLimit));

        return output.sub(outputAddress).toInt();
    }

    private static compressFrame(inputBase: any, inputAddress: Long, inputLimit: Long, outputBase: any, outputAddress: Long, outputLimit: Long, parameters: CompressionParameters): number {
        let windowSize: number = 1 << parameters.getWindowLog(); // TODO: store window size in parameters directly?
        let blockSize: number = Math.min(Constants.MAX_BLOCK_SIZE, windowSize);

        let outputSize: number = outputLimit.sub(outputAddress).toInt();
        let remaining: number = inputLimit.sub(inputAddress).toInt();

        let output: Long = outputAddress;
        let input: Long = inputAddress;

        let context: CompressionContext = new CompressionContext(parameters, inputAddress, remaining);

        do {
            Util.checkArgument(outputSize >= Constants.SIZE_OF_BLOCK_HEADER + Constants.MIN_BLOCK_SIZE, "Output buffer too small");

            let lastBlockFlag: number = blockSize >= remaining ? 1 : 0;
            blockSize = Math.min(blockSize, remaining);

            let compressedSize: number = 0;
            if (remaining > 0) {
                compressedSize = ZstdFrameCompressor.compressBlock(inputBase, input, blockSize, outputBase, output.add(Constants.SIZE_OF_BLOCK_HEADER), outputSize - Constants.SIZE_OF_BLOCK_HEADER, context, parameters);
            }

            if (compressedSize == 0) { // block is not compressible
                Util.checkArgument(blockSize + Constants.SIZE_OF_BLOCK_HEADER <= outputSize, "Output size too small");

                let blockHeader: number = lastBlockFlag | (Constants.RAW_BLOCK << 1) | (blockSize << 3);
                Util.put24BitLittleEndian(outputBase, output, blockHeader);
                Unsafe.copyMemory(inputBase, input.toNumber(), outputBase,
                output.add(Constants.SIZE_OF_BLOCK_HEADER).toNumber(), Long.fromNumber(blockSize));
                compressedSize = Constants.SIZE_OF_BLOCK_HEADER + blockSize;
            } else {
                let blockHeader: number = lastBlockFlag | (Constants.COMPRESSED_BLOCK << 1) | (compressedSize << 3);
                Util.put24BitLittleEndian(outputBase, output, blockHeader);
                compressedSize += Constants.SIZE_OF_BLOCK_HEADER;
            }

            input = input.add(blockSize);
            remaining -= blockSize;
            output = output.add(compressedSize);
            outputSize -= compressedSize;
        }
        while (remaining > 0);

        return output.sub(outputAddress).toInt();
    }

    private static compressBlock(inputBase: any, inputAddress: Long, inputSize: number, outputBase: any, outputAddress: Long, outputSize: number, context: CompressionContext, parameters: CompressionParameters): number
    {
        if (inputSize < Constants.MIN_BLOCK_SIZE + Constants.SIZE_OF_BLOCK_HEADER + 1) {
            return 0;
        }

        context.blockCompressionState.enforceMaxDistance(inputAddress.add(inputSize), 1 << parameters.getWindowLog());
        context.sequenceStore.reset();

        let lastLiteralsSize: number = parameters.getStrategy()
            .getCompressor()
            .compressBlock(inputBase, inputAddress, inputSize, context.sequenceStore, context.blockCompressionState, context.offsets, parameters);

        let lastLiteralsAddress: Long = inputAddress.add(inputSize).sub(lastLiteralsSize);

        context.sequenceStore.appendLiterals(inputBase, lastLiteralsAddress, lastLiteralsSize);

        context.sequenceStore.generateCodes();

        let outputLimit: Long = outputAddress.add(outputSize);
        let output: Long = outputAddress;

        let compressedLiteralsSize: number = ZstdFrameCompressor.encodeLiterals(
            context.huffmanContext,
            parameters,
            outputBase,
            output,
        outputLimit.sub(output).toInt(),
            context.sequenceStore.literalsBuffer,
            context.sequenceStore.literalsLength);
        output = output.add(compressedLiteralsSize);

        let compressedSequencesSize: number = SequenceEncoder.compressSequences(outputBase, output,
        outputLimit.sub(output)
            .toInt(), context.sequenceStore, parameters.getStrategy(), context.sequenceEncodingContext);

        let compressedSize: number = compressedLiteralsSize + compressedSequencesSize;
        if (compressedSize == 0) {
            return compressedSize;
        }

        let maxCompressedSize: number = inputSize - ZstdFrameCompressor.calculateMinimumGain(inputSize, parameters.getStrategy());
        if (compressedSize > maxCompressedSize) {
            return 0; // not compressed
        }

        context.commit();

        return compressedSize;
    }

    private static encodeLiterals(
        context: HuffmanCompressionContext,
        parameters: CompressionParameters,
        outputBase: any,
        outputAddress: Long,
        outputSize: number,
        literals: Int8Array,
        literalsSize: number): number
    {
        let bypassCompression: boolean = (parameters.getStrategy() == Strategy.FAST) && (parameters.getTargetLength() > 0);
        if (bypassCompression || literalsSize <= ZstdFrameCompressor.MINIMUM_LITERALS_SIZE) {
            return ZstdFrameCompressor.rawLiterals(outputBase, outputAddress, outputSize, literals, Long.fromNumber(Unsafe.ARRAY_BYTE_BASE_OFFSET), literalsSize);
        }

        let headerSize: number = 3 + (literalsSize >= 1024 ? 1 : 0) + (literalsSize >= 16384 ? 1 : 0);

        Util.checkArgument(headerSize + 1 <= outputSize, "Output buffer too small");

        let counts: Int32Array = new Int32Array(Huffman.MAX_SYMBOL_COUNT);
        Histogram.count(literals, literalsSize, counts);
        let maxSymbol: number = Histogram.findMaxSymbol(counts, Huffman.MAX_SYMBOL);
        let largestCount: number = Histogram.findLargestCount(counts, maxSymbol);

        let literalsAddress: Long = Long.fromNumber(Unsafe.ARRAY_BYTE_BASE_OFFSET);
        if (largestCount == literalsSize) {
            // all bytes in input are equal
            return ZstdFrameCompressor.rleLiterals(outputBase, outputAddress, outputSize, literals, Long.fromNumber(Unsafe.ARRAY_BYTE_BASE_OFFSET), literalsSize);
        }
        else if (largestCount <= (literalsSize >>> 7) + 4) {
            // heuristic: probably not compressible enough
            return ZstdFrameCompressor.rawLiterals(outputBase, outputAddress, outputSize, literals, Long.fromNumber(Unsafe.ARRAY_BYTE_BASE_OFFSET), literalsSize);
        }

        let previousTable: HuffmanCompressionTable = context.getPreviousTable();
        let table: HuffmanCompressionTable;
        let serializedTableSize: number;
        let reuseTable: boolean;

        let canReuse: boolean = previousTable.isValid(counts, maxSymbol);

        let preferReuse: boolean = parameters.getStrategy().ordinal() < Strategy.LAZY.ordinal() && literalsSize <= 1024;
        if (preferReuse && canReuse) {
            table = previousTable;
            reuseTable = true;
            serializedTableSize = 0;
        }
        else {
            let newTable: HuffmanCompressionTable = context.borrowTemporaryTable();

            newTable.initialize(
                counts,
                maxSymbol,
            HuffmanCompressionTable.optimalNumberOfBits(ZstdFrameCompressor.MAX_HUFFMAN_TABLE_LOG, literalsSize, maxSymbol),
            context.getCompressionTableWorkspace());

            serializedTableSize = newTable.write(outputBase, outputAddress.add(headerSize), outputSize - headerSize, context.getTableWriterWorkspace());

            if (canReuse && previousTable.estimateCompressedSize(counts, maxSymbol) <= serializedTableSize + newTable.estimateCompressedSize(counts, maxSymbol)) {
                table = previousTable;
                reuseTable = true;
                serializedTableSize = 0;
                context.discardTemporaryTable();
            }
            else {
                table = newTable;
                reuseTable = false;
            }
        }

        let compressedSize: number;
        let singleStream: boolean = literalsSize < 256;
        if (singleStream) {
            compressedSize = HuffmanCompressor.compressSingleStream(outputBase,
            outputAddress.add(headerSize)
                .add(serializedTableSize), outputSize - headerSize - serializedTableSize, literals, literalsAddress, literalsSize, table);
        }
        else {
            compressedSize = HuffmanCompressor.compress4streams(outputBase,
            outputAddress.add(headerSize)
                .add(serializedTableSize), outputSize - headerSize - serializedTableSize, literals, literalsAddress, literalsSize, table);
        }

        let totalSize: number = serializedTableSize + compressedSize;
        let minimumGain: number = ZstdFrameCompressor.calculateMinimumGain(literalsSize, parameters.getStrategy());

        if (compressedSize == 0 || totalSize >= literalsSize - minimumGain) {

            context.discardTemporaryTable();

            return ZstdFrameCompressor.rawLiterals(outputBase, outputAddress, outputSize, literals, Long.fromNumber(Unsafe.ARRAY_BYTE_BASE_OFFSET), literalsSize);
        }

        let encodingType: number = reuseTable ? Constants.TREELESS_LITERALS_BLOCK : Constants.COMPRESSED_LITERALS_BLOCK;

        switch (headerSize) {
            case 3:
            { // 2 - 2 - 10 - 10
                let header: number = encodingType | ((singleStream ? 0 : 1) << 2) | (literalsSize << 4) | (totalSize << 14);
                Util.put24BitLittleEndian(outputBase, outputAddress, header);
                break;
            }
            case 4:
            { // 2 - 2 - 14 - 14
                let header: number = encodingType | (2 << 2) | (literalsSize << 4) | (totalSize << 18);
                Unsafe.putInt(outputBase, outputAddress.toNumber(), header);
                break;
            }
            case 5:
            { // 2 - 2 - 18 - 18
                let header: number = encodingType | (3 << 2) | (literalsSize << 4) | (totalSize << 22);
                Unsafe.putInt(outputBase, outputAddress.toNumber(), header);
                Unsafe.putByte(outputBase,
                outputAddress.add(Constants.SIZE_OF_INT).toNumber(), NumberTransform.toByte(totalSize >>> 10));
                break;
            }
            default:
                throw new Exception();
        }

        return headerSize + totalSize;
    }

    private static rleLiterals(outputBase: any, outputAddress: Long, outputSize: number, inputBase: any, inputAddress: Long, inputSize: number): number {
        let headerSize: number = 1 + (inputSize > 31 ? 1 : 0) + (inputSize > 4095 ? 1 : 0);

        switch (headerSize) {
            case 1: // 2 - 1 - 5
                Unsafe.putByte(outputBase, outputAddress.toNumber(), NumberTransform.toByte(Constants.RLE_LITERALS_BLOCK) | (inputSize << 3));
                break;
            case 2: // 2 - 2 - 12
                Unsafe.putShort(outputBase, outputAddress.toNumber(), NumberTransform.toShort(Constants.RLE_LITERALS_BLOCK) | (1 << 2) | (inputSize << 4));
                break;
            case 3: // 2 - 2 - 20
                Unsafe.putInt(outputBase, outputAddress.toNumber(), Constants.RLE_LITERALS_BLOCK | 3 << 2 | inputSize << 4);
                break;
            default: // impossible. headerSize is {1,2,3}
                throw new Exception();
        }

        Unsafe.putByte(outputBase,
        outputAddress.add(headerSize).toNumber(), Unsafe.getByte(inputBase, inputAddress.toNumber()));

        return headerSize + 1;
    }

    private static calculateMinimumGain(inputSize: number, strategy: Strategy): number{
        // TODO: move this to Strategy to avoid hardcoding a specific strategy here
        let minLog: number = strategy == Strategy.BTULTRA ? 7 : 6;
        return (inputSize >>> minLog) + 2;
    }

    private static rawLiterals(outputBase: any, outputAddress: Long, outputSize: number, inputBase: any, inputAddress: Long, inputSize: number): number
    {
        let headerSize: number = 1;
        if (inputSize >= 32) {
            headerSize++;
        }
        if (inputSize >= 4096) {
            headerSize++;
        }

        Util.checkArgument(inputSize + headerSize <= outputSize, "Output buffer too small");

        switch (headerSize) {
            case 1:
                Unsafe.putByte(outputBase, outputAddress.toNumber(), NumberTransform.toByte(Constants.RAW_LITERALS_BLOCK | (inputSize << 3)));
                break;
            case 2:
                Unsafe.putShort(outputBase, outputAddress.toNumber(), NumberTransform.toShort(Constants.RAW_LITERALS_BLOCK | (1 << 2) | (inputSize << 4)));
                break;
            case 3:
                Util.put24BitLittleEndian(outputBase, outputAddress, Constants.RAW_LITERALS_BLOCK | (3 << 2) | (inputSize << 4));
                break;
            default:
                throw new Error();
        }

        Util.checkArgument(inputSize + 1 <= outputSize, "Output buffer too small");

        Unsafe.copyMemory(inputBase, inputAddress.toNumber(), outputBase,
        outputAddress.add(headerSize).toNumber(), Long.fromNumber(inputSize));

        return headerSize + inputSize;
    }
}
